#include "lockfree_hashtable.h"
#include "../atomic/cpu_pause.h"
#include <stdlib.h>
#include <string.h>

bool hashtable_init(hashtable_t* table, 
                   size_t bucket_count,
                   size_t (*hash_func)(const void* key),
                   bool (*key_equals)(const void* key1, const void* key2)) {
    
    table->buckets = calloc(bucket_count, sizeof(hash_bucket_t));
    if (!table->buckets) return false;
    
    table->bucket_count = bucket_count;
    table->hash_func = hash_func;
    table->key_equals = key_equals;
    
    // 初始化所有桶
    for (size_t i = 0; i < bucket_count; i++) {
        tagged_ptr_t initial = make_tagged_ptr(NULL, 0);
        atomic_tagged_store(&table->buckets[i].head, initial, memory_order_relaxed);
    }
    
    return true;
}

void hashtable_destroy(hashtable_t* table) {
    // 清理所有桶中的节点
    for (size_t i = 0; i < table->bucket_count; i++) {
        hash_node_t* current = get_ptr(
            atomic_tagged_load(&table->buckets[i].head, memory_order_relaxed)
        );
        while (current) {
            hash_node_t* next = current->next;
            free(current);
            current = next;
        }
    }
    free(table->buckets);
}

static hash_bucket_t* get_bucket(hashtable_t* table, const void* key) {
    size_t hash = table->hash_func(key);
    return &table->buckets[hash % table->bucket_count];
}

bool hashtable_put(hashtable_t* table, void* key, void* value) {
    hash_bucket_t* bucket = get_bucket(table, key);
    size_t hash = table->hash_func(key);
    
    // 创建新节点
    hash_node_t* new_node = malloc(sizeof(hash_node_t));
    if (!new_node) return false;
    
    new_node->key = key;
    new_node->value = value;
    new_node->hash = hash;
    
    do {
        tagged_ptr_t head = atomic_tagged_load(&bucket->head, memory_order_acquire);
        hash_node_t* current = get_ptr(head);
        
        // 检查是否已存在相同的键
        while (current) {
            if (current->hash == hash && table->key_equals(current->key, key)) {
                free(new_node);
                return false;  // 键已存在
            }
            current = current->next;
        }
        
        // 将新节点插入到链表头部
        new_node->next = get_ptr(head);
        tagged_ptr_t new_head = make_tagged_ptr(new_node, get_tag(head) + 1);
        
        if (atomic_tagged_compare_exchange_weak(
                &bucket->head,
                &head,
                new_head,
                memory_order_release,
                memory_order_relaxed)) {
            return true;
        }
        cpu_pause();
    } while (1);
}

void* hashtable_get(hashtable_t* table, const void* key) {
    hash_bucket_t* bucket = get_bucket(table, key);
    size_t hash = table->hash_func(key);
    
    tagged_ptr_t head = atomic_tagged_load(&bucket->head, memory_order_acquire);
    hash_node_t* current = get_ptr(head);
    
    while (current) {
        if (current->hash == hash && table->key_equals(current->key, key)) {
            return current->value;
        }
        current = current->next;
    }
    
    return NULL;
}

bool hashtable_remove(hashtable_t* table, const void* key) {
    hash_bucket_t* bucket = get_bucket(table, key);
    size_t hash = table->hash_func(key);
    
    do {
        tagged_ptr_t head = atomic_tagged_load(&bucket->head, memory_order_acquire);
        hash_node_t* current = get_ptr(head);
        hash_node_t* prev = NULL;
        
        // 查找要删除的节点
        while (current) {
            if (current->hash == hash && table->key_equals(current->key, key)) {
                if (!prev) {
                    // 删除头节点
                    tagged_ptr_t new_head = make_tagged_ptr(
                        current->next,
                        get_tag(head) + 1
                    );
                    if (atomic_tagged_compare_exchange_weak(
                            &bucket->head,
                            &head,
                            new_head,
                            memory_order_release,
                            memory_order_relaxed)) {
                        free(current);
                        return true;
                    }
                } else {
                    // 删除中间节点
                    if (__sync_bool_compare_and_swap(
                            &prev->next,
                            current,
                            current->next)) {
                        free(current);
                        return true;
                    }
                }
                break;
            }
            prev = current;
            current = current->next;
        }
        
        if (!current) return false;  // 未找到要删除的键
        cpu_pause();
    } while (1);
}

size_t hashtable_size(hashtable_t* table) {
    size_t count = 0;
    for (size_t i = 0; i < table->bucket_count; i++) {
        tagged_ptr_t head = atomic_tagged_load(
            &table->buckets[i].head,
            memory_order_acquire
        );
        hash_node_t* current = get_ptr(head);
        while (current) {
            count++;
            current = current->next;
        }
    }
    return count;
} 